//! 技术分析工具模块
//! 
//! 提供各种技术分析指标的计算工具

use serde::{Deserialize, Serialize};
use crate::objects::RawBar;

/// 技术分析工具配置
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaConfig {
    /// 移动平均周期
    pub ma_periods: Vec<usize>,
    /// MACD参数
    pub macd_fast: usize,
    pub macd_slow: usize,
    pub macd_signal: usize,
    /// RSI周期
    pub rsi_period: usize,
    /// 布林带周期
    pub bb_period: usize,
    pub bb_std: f64,
}

/// 技术分析工具
#[derive(Debug, Clone)]
pub struct TaUtils {
    config: TaConfig,
}

impl TaUtils {
    /// 创建新的技术分析工具
    pub fn new(config: TaConfig) -> Self {
        Self { config }
    }

    /// 计算简单移动平均
    pub fn sma(&self, prices: &[f64], period: usize) -> Vec<f64> {
        if prices.len() < period {
            return vec![];
        }
        
        let mut ma_values = Vec::new();
        for i in period - 1..prices.len() {
            let start_idx = if i >= period - 1 { i.saturating_sub(period - 1) } else { 0 };
            let sum: f64 = prices[start_idx..=i].iter().sum();
            ma_values.push(sum / period as f64);
        }
        ma_values
    }

    /// 计算指数移动平均
    pub fn ema(&self, prices: &[f64], period: usize) -> Vec<f64> {
        if prices.is_empty() {
            return vec![];
        }
        
        let alpha = 2.0 / (period + 1) as f64;
        let mut ema_values = Vec::new();
        let mut ema = prices[0];
        ema_values.push(ema);
        
        for &price in &prices[1..] {
            ema = alpha * price + (1.0 - alpha) * ema;
            ema_values.push(ema);
        }
        
        ema_values
    }

    /// 计算加权移动平均
    pub fn wma(&self, prices: &[f64], period: usize) -> Vec<f64> {
        if prices.len() < period {
            return vec![];
        }
        
        let mut wma_values = Vec::new();
        for i in period - 1..prices.len() {
            let mut weighted_sum = 0.0;
            let mut weight_sum = 0.0;
            
            for (j, &price) in prices[i - period + 1..=i].iter().enumerate() {
                let weight = (j + 1) as f64;
                weighted_sum += price * weight;
                weight_sum += weight;
            }
            
            wma_values.push(weighted_sum / weight_sum);
        }
        wma_values
    }

    /// 计算MACD
    pub fn macd(&self, prices: &[f64]) -> (Vec<f64>, Vec<f64>, Vec<f64>) {
        let fast_ema = self.ema(prices, self.config.macd_fast);
        let slow_ema = self.ema(prices, self.config.macd_slow);
        
        if fast_ema.len() != slow_ema.len() {
            return (vec![], vec![], vec![]);
        }
        
        let mut macd_line = Vec::new();
        for i in 0..fast_ema.len() {
            macd_line.push(fast_ema[i] - slow_ema[i]);
        }
        
        let signal_line = self.ema(&macd_line, self.config.macd_signal);
        let mut histogram = Vec::new();
        
        for i in 0..macd_line.len() {
            if i < signal_line.len() {
                histogram.push(macd_line[i] - signal_line[i]);
            } else {
                histogram.push(0.0);
            }
        }
        
        (macd_line, signal_line, histogram)
    }

    /// 计算RSI
    pub fn rsi(&self, prices: &[f64]) -> Vec<f64> {
        if prices.len() < self.config.rsi_period + 1 {
            return vec![];
        }
        
        let mut gains = Vec::new();
        let mut losses = Vec::new();
        
        for i in 1..prices.len() {
            let change = prices[i] - prices[i-1];
            gains.push(change.max(0.0));
            losses.push((-change).max(0.0));
        }
        
        let avg_gain = self.ema(&gains, self.config.rsi_period);
        let avg_loss = self.ema(&losses, self.config.rsi_period);
        
        let mut rsi_values = Vec::new();
        for i in 0..avg_gain.len() {
            if avg_loss[i] == 0.0 {
                rsi_values.push(100.0);
            } else {
                let rs = avg_gain[i] / avg_loss[i];
                let rsi = 100.0 - (100.0 / (1.0 + rs));
                rsi_values.push(rsi);
            }
        }
        
        rsi_values
    }

    /// 计算布林带
    pub fn bollinger_bands(&self, prices: &[f64]) -> (Vec<f64>, Vec<f64>, Vec<f64>) {
        let sma = self.sma(prices, self.config.bb_period);
        let mut upper_band = Vec::new();
        let mut lower_band = Vec::new();
        
        for i in 0..sma.len() {
            let start_idx = i;
            let end_idx = (i + self.config.bb_period).min(prices.len());
            if start_idx >= end_idx {
                continue;
            }
            let slice = &prices[start_idx..end_idx];
            
            let mean = slice.iter().sum::<f64>() / slice.len() as f64;
            let variance: f64 = slice.iter()
                .map(|x| (x - mean).powi(2))
                .sum::<f64>() / slice.len() as f64;
            let std_dev = variance.sqrt();
            
            upper_band.push(mean + self.config.bb_std * std_dev);
            lower_band.push(mean - self.config.bb_std * std_dev);
        }
        
        (upper_band, sma, lower_band)
    }

    /// 计算随机指标
    pub fn stochastic(&self, highs: &[f64], lows: &[f64], closes: &[f64], period: usize) -> (Vec<f64>, Vec<f64>) {
        if highs.len() < period || lows.len() < period || closes.len() < period {
            return (vec![], vec![]);
        }
        
        let mut k_values = Vec::new();
        let mut d_values = Vec::new();
        
        for i in period - 1..closes.len() {
            let high_slice = &highs[i - period + 1..=i];
            let low_slice = &lows[i - period + 1..=i];
            
            let highest_high = high_slice.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b));
            let lowest_low = low_slice.iter().fold(f64::INFINITY, |a, &b| a.min(b));
            
            let k = if highest_high == lowest_low {
                50.0
            } else {
                ((closes[i] - lowest_low) / (highest_high - lowest_low)) * 100.0
            };
            k_values.push(k);
        }
        
        // 计算%D (K的移动平均)
        if k_values.len() >= 3 {
            d_values = self.sma(&k_values, 3);
        }
        
        (k_values, d_values)
    }

    /// 计算威廉指标
    pub fn williams_r(&self, highs: &[f64], lows: &[f64], closes: &[f64], period: usize) -> Vec<f64> {
        if highs.len() < period || lows.len() < period || closes.len() < period {
            return vec![];
        }
        
        let mut wr_values = Vec::new();
        
        for i in period - 1..closes.len() {
            let high_slice = &highs[i - period + 1..=i];
            let low_slice = &lows[i - period + 1..=i];
            
            let highest_high = high_slice.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b));
            let lowest_low = low_slice.iter().fold(f64::INFINITY, |a, &b| a.min(b));
            
            let wr = if highest_high == lowest_low {
                -50.0
            } else {
                ((highest_high - closes[i]) / (highest_high - lowest_low)) * -100.0
            };
            wr_values.push(wr);
        }
        
        wr_values
    }

    /// 计算ATR (Average True Range)
    pub fn atr(&self, highs: &[f64], lows: &[f64], closes: &[f64], period: usize) -> Vec<f64> {
        if highs.len() < 2 || lows.len() < 2 || closes.len() < 2 {
            return vec![];
        }
        
        let mut true_ranges = Vec::new();
        
        for i in 1..closes.len() {
            let high_low = highs[i] - lows[i];
            let high_close = (highs[i] - closes[i-1]).abs();
            let low_close = (lows[i] - closes[i-1]).abs();
            
            let true_range = high_low.max(high_close).max(low_close);
            true_ranges.push(true_range);
        }
        
        self.sma(&true_ranges, period)
    }

    /// 计算CCI (Commodity Channel Index)
    pub fn cci(&self, highs: &[f64], lows: &[f64], closes: &[f64], period: usize) -> Vec<f64> {
        if highs.len() < period || lows.len() < period || closes.len() < period {
            return vec![];
        }
        
        let mut cci_values = Vec::new();
        
        for i in period - 1..closes.len() {
            let typical_prices: Vec<f64> = (i - period + 1..=i)
                .map(|j| (highs[j] + lows[j] + closes[j]) / 3.0)
                .collect();
            
            let sma = typical_prices.iter().sum::<f64>() / period as f64;
            let mean_deviation: f64 = typical_prices.iter()
                .map(|&tp| (tp - sma).abs())
                .sum::<f64>() / period as f64;
            
            let cci = if mean_deviation != 0.0 {
                (typical_prices.last().unwrap() - sma) / (0.015 * mean_deviation)
            } else {
                0.0
            };
            cci_values.push(cci);
        }
        
        cci_values
    }

    /// 计算所有技术指标
    pub fn calculate_all_indicators(&self, bars: &[RawBar]) -> serde_json::Value {
        if bars.is_empty() {
            return serde_json::json!({
                "error": "K线数据为空，无法计算指标"
            });
        }
        
        let prices: Vec<f64> = bars.iter().map(|bar| bar.close).collect();
        let highs: Vec<f64> = bars.iter().map(|bar| bar.high).collect();
        let lows: Vec<f64> = bars.iter().map(|bar| bar.low).collect();
        
        // 计算各种技术指标
        let ma5 = self.sma(&prices, 5);
        let ma10 = self.sma(&prices, 10);
        let ma20 = self.sma(&prices, 20);
        let ema12 = self.ema(&prices, 12);
        let ema26 = self.ema(&prices, 26);
        
        let (macd_line, signal_line, histogram) = self.macd(&prices);
        let rsi = self.rsi(&prices);
        let (bb_upper, bb_middle, bb_lower) = self.bollinger_bands(&prices);
        let (stoch_k, stoch_d) = self.stochastic(&highs, &lows, &prices, 14);
        let williams_r = self.williams_r(&highs, &lows, &prices, 14);
        let atr = self.atr(&highs, &lows, &prices, 14);
        let cci = self.cci(&highs, &lows, &prices, 14);
        
        serde_json::json!({
            "moving_averages": {
                "sma_5": ma5,
                "sma_10": ma10,
                "sma_20": ma20,
                "ema_12": ema12,
                "ema_26": ema26
            },
            "macd": {
                "line": macd_line,
                "signal": signal_line,
                "histogram": histogram
            },
            "rsi": rsi,
            "bollinger_bands": {
                "upper": bb_upper,
                "middle": bb_middle,
                "lower": bb_lower
            },
            "stochastic": {
                "k": stoch_k,
                "d": stoch_d
            },
            "williams_r": williams_r,
            "atr": atr,
            "cci": cci,
            "price_count": prices.len(),
            "indicator_count": {
                "ma": ma5.len(),
                "macd": macd_line.len(),
                "rsi": rsi.len(),
                "bb": bb_upper.len(),
                "stoch": stoch_k.len(),
                "atr": atr.len(),
                "cci": cci.len()
            }
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_sma() {
        let config = TaConfig {
            ma_periods: vec![5, 10, 20],
            macd_fast: 12,
            macd_slow: 26,
            macd_signal: 9,
            rsi_period: 14,
            bb_period: 20,
            bb_std: 2.0,
        };
        let ta_utils = TaUtils::new(config);
        
        let prices = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0];
        let ma = ta_utils.sma(&prices, 5);
        
        assert_eq!(ma.len(), 6);
        assert_eq!(ma[0], 3.0); // (1+2+3+4+5)/5
        assert_eq!(ma[1], 4.0); // (2+3+4+5+6)/5
    }

    #[test]
    fn test_ema() {
        let config = TaConfig {
            ma_periods: vec![5, 10, 20],
            macd_fast: 12,
            macd_slow: 26,
            macd_signal: 9,
            rsi_period: 14,
            bb_period: 20,
            bb_std: 2.0,
        };
        let ta_utils = TaUtils::new(config);
        
        let prices = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let ema = ta_utils.ema(&prices, 3);
        
        assert_eq!(ema.len(), 5);
        assert_eq!(ema[0], 1.0); // 初始值
    }

    #[test]
    fn test_rsi() {
        let config = TaConfig {
            ma_periods: vec![5, 10, 20],
            macd_fast: 12,
            macd_slow: 26,
            macd_signal: 9,
            rsi_period: 14,
            bb_period: 20,
            bb_std: 2.0,
        };
        let ta_utils = TaUtils::new(config);
        
        let prices = vec![44.0, 44.34, 44.09, 44.15, 43.61, 44.33, 44.34, 44.09, 44.15, 43.61];
        let rsi = ta_utils.rsi(&prices);
        
        assert!(!rsi.is_empty());
        assert!(rsi.iter().all(|&x| x >= 0.0 && x <= 100.0));
    }
} 